import random
import numpy as np
import time
import math
import matplotlib.pyplot as plt

class Plot:
    @staticmethod
    def reset():
        plt.cla()
        plt.axis("equal")
        plt.grid(True)

    @staticmethod
    def plot_data(m, color="ob"):
        plt.plot(m[0, :], m[1, :], color, markersize=2)

'''
简易的搜索匹配点云方法
'''

# ==========真实位置=========== #

def get_ob(max_sample=50, x_len=3.0, width=2.0, length=1.0, noisy=False, fake_point=False, noisy_seed=0.05):
    horizon_wall = []
    vertical_wall = []
    
    def rd(seed):
        return random.uniform(-1,1)*seed

    if not noisy: noisy_seed = 0
    for i in range(max_sample):
        horizon_wall.append([-x_len-width+i*width/max_sample+rd(noisy_seed), 0+rd(noisy_seed), 1.0])
        vertical_wall.append([-x_len+rd(noisy_seed), i*length/max_sample+rd(noisy_seed), 1.0])

    _world = np.vstack(
        [np.array(horizon_wall),
        np.array(vertical_wall)]
        )

    mirror = np.vstack([
        -1*_world[:, 0],
        _world[:, 1],
        _world[:, 2]
    ]).T

    _world = np.vstack(
        [_world, mirror]
    )

    _fake = []
    for i in range(20):
        _fake.append([rd(1.0), 3.0, 1.0])

    if fake_point:
        _world = np.vstack([
            _world, np.array(_fake)
        ])

    return _world.T


# ========采样值========== #

def get_trans_matrix(dx, dy):
    return np.array([
        [1.0, 0.0, dx],
        [0.0, 1.0, dy],
        [0.0, 0.0, 1]
    ])

def get_rot_matrix(dth):
    return np.array([
        [math.cos(dth), -math.sin(dth), 0.0],
        [math.sin(dth), math.cos(dth), 0.0],
        [0.0, 0.0, 1.0]
    ])

score_cnt = 0
def match_score(scan, real):
    '''
    得到真实real世界和扫描scan世界的匹配系数
    '''
    def single_score(s):
        dis = []
        global score_cnt
        dis = map(lambda r:(math.sqrt((r[0]-s[0])**2 + (r[1]-s[1])**2)), (real.T))
        ret = min(dis)
        return ret
    
    score = map(lambda s:single_score(s), scan.T)
    return sum(score)/len(scan)

def icp(scan, real, x_step=0.1, y_step=0.1, th_step=3):
    best_x = 0.0
    best_y = 0.0
    best_th = 0.0

    new_scan = None
    new_score = None
    origin_scan = scan
    origin_score = match_score(origin_scan, real)

    def one_x_step(step):
        nonlocal new_scan, new_score
        new_scan = origin_scan + [[step], [0.0], [0.0]]
        new_score = match_score(new_scan, real)

    def one_y_step(step):
        nonlocal new_scan, new_score
        new_scan = origin_scan + [[0.0], [step], [0.0]]
        new_score = match_score(new_scan, real)

    def one_th_step(step):
        step = math.radians(step)
        nonlocal new_scan, new_score
        trans_matrix = np.array([
            [math.cos(step), -math.sin(step), 0.0],
            [math.sin(step), math.cos(step), 0.0],
            [0.0, 0.0, 1.0]
        ])
        new_scan = trans_matrix.dot(origin_scan)
        new_score = match_score(new_scan, real)

    def explore(func, step):
        nonlocal origin_score, origin_scan
        best = 0
        # 先试探
        func(step)
        # 说明方向错了
        if new_score>origin_score: 
            step = -step
            func(step)
        best += step
        # 变步长
        auto_step = step
        # 正式开始
        while new_score < origin_score:
            origin_score = new_score
            origin_scan = new_scan
            func(auto_step)
            auto_step = auto_step * 1.1
            best += auto_step
        # 去掉最后一次不符合条件多加的
        best = best - auto_step
        return best

    best_x = explore(one_x_step, x_step)
    print("x", best_x)
    best_y = explore(one_y_step, y_step)
    print("y", best_y)
    best_th = explore(one_th_step, th_step)
    print(best_x, best_y, math.radians(best_th))
    return (best_x, best_y, math.radians(best_th))

def run_demo():
    print("!! start !!")
    real_world = get_ob(20)
    rot_angle = math.pi/3
    rot_matrix = np.array([
            [math.cos(rot_angle), -math.sin(rot_angle), 0.0],
            [math.sin(rot_angle), math.cos(rot_angle), 0.0],
            [0.0, 0.0, 1.0]
        ])

    trans = [-1.0, 2.0]
    trans_matrix = np.array([
        [1.0, 0.0, trans[0]],
        [0.0, 1.0, trans[1]],
        [0.0, 0.0, 1]
    ])
    laser_scan = trans_matrix.dot(
                        rot_matrix.dot(
                            get_ob(20, noisy=False, fake_point=False)))

    origin_scan = laser_scan

    Plot.reset()
    Plot.plot_data(real_world, 'or')
    Plot.plot_data(origin_scan)
    plt.pause(0.001)

    best_matrix = np.array([
        [1, 0, 0],
        [0, 1, 0],
        [0, 0, 1]
    ])

    score_cache_0 = 10000.0
    score_cache_1 = 0.0
    score_threshold = 100.0

    start_time = time.time()
    print(" Start Loop ")
    while abs(score_cache_0-score_cache_1)>0.1 or abs(score_cache_1/len(laser_scan[0]))>0.5 or True:

        if time.time() - start_time > 10:
            print('[Time Out]')
            break

        score_cache_0 = score_cache_1
        score_cache_1 = match_score(origin_scan, real_world)

        dx, dy, dth = icp(origin_scan, real_world, 0.01, 0.01, 0.1)
        T = get_trans_matrix(dx, dy)
        R = get_rot_matrix(dth)
        best_matrix = R.dot(T.dot(best_matrix))

        # 计算平移距离和旋转角度
        xx = (best_matrix[0,2]/best_matrix[1,0]+best_matrix[1,2]/best_matrix[0,0])/(best_matrix[0,0]/best_matrix[1,0]+best_matrix[1,0]/best_matrix[0,0])
        yy = best_matrix[1,2]/best_matrix[0,0]-xx*best_matrix[1,0]/best_matrix[0,0]
        tth = math.degrees(math.acos(best_matrix[0,0]))
        print("Transform: [%.3f, %.3f, %.2f]"%(xx, yy, tth))

        best_match = R.dot(T.dot(origin_scan))
        Plot.reset()
        Plot.plot_data(real_world, 'or')
        Plot.plot_data(laser_scan)
        Plot.plot_data(best_matrix.dot(laser_scan), 'og')
        plt.pause(0.001)
        origin_scan = best_match

        print(match_score(best_match, real_world))

    print('Score Cnt', score_cnt)
    Plot.reset()
    Plot.plot_data(real_world, 'or')
    Plot.plot_data(laser_scan)
    Plot.plot_data(best_match, 'og')
    plt.show()

if __name__ == "__main__":
    run_demo()

